import numpy as np
import cv2
import os
from os.path import join as ospj

from paddle import fluid

from utils.logger import get_logger
logger = get_logger(__name__)


def translate_using_reference(nets, args, x_src, x_ref, y_ref, filename):
    N, C, H, W = x_src.shape
    wb = fluid.dygraph.to_variable(np.ones((1, C, H, W))).astype('float32')
    x_src_with_wb = fluid.layers.concat([wb, x_src], axis=0)

    masks = nets.fan.get_heatmap(x_src) if args.w_hpf > 0 else None
    s_ref = nets.style_encoder(x_ref, y_ref)
    s_ref_list = fluid.layers.unsqueeze(s_ref, axes=[1])
    s_ref_lists = []
    for _ in range(N):
        s_ref_lists.append(s_ref_list)
    s_ref_list = fluid.layers.stack(s_ref_lists, axis=1)
    s_ref_list = fluid.layers.reshape(s_ref_list, (s_ref_list.shape[0], s_ref_list.shape[1], s_ref_list.shape[3]))
    x_concat = [x_src_with_wb]
    for i, s_ref in enumerate(s_ref_list):
        x_fake = nets.generator(x_src, s_ref, masks=masks)
        x_fake_with_ref = fluid.layers.concat([x_ref[i:i+1], x_fake], axis=0)
        x_concat += [x_fake_with_ref]

    x_concat = fluid.layers.concat(x_concat, axis=0)
    save_images(x_concat, N+1, filename, rotate=True)
    del x_concat


def denormalize(x):
    out = (x + 1) / 2
    return fluid.layers.clip(out, 0, 1)


def save_images(x, ncol, filename, rotate=False):
    x = denormalize(x)
    save_image(x, filename, nrow=ncol, padding=0, rotate=rotate)


def make_grid(tensor, nrow=8, padding=2,
              normalize=False, range=None, scale_each=False, pad_value=0, rotate=False):
    images = tensor.numpy()
    n, c, h, w = images.shape
    ncol = int(n / nrow)
    if rotate:
        nrow, ncol = ncol, nrow
    dw = w * ncol
    dh = h * nrow
    dst_img = np.ones((dh, dw, c), dtype=np.uint8)
    for i, img in enumerate(images):
        img = np.transpose(img, (1, 2, 0))
        img = img * 255
        img = img + 0.5
        img = np.clip(img, 0, 255).astype(np.uint8)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        if not rotate:
            row = i % nrow
            col = int(i / nrow)
        else:
            col = i % ncol
            row = int(i / ncol)
        sw = col * w
        dw = sw + w
        sh = row * h
        dh = sh + h
        dst_img[sh:dh, sw:dw] = img
    return dst_img


def save_image(tensor, fp, nrow=8, padding=2,
               normalize=False, range=None, scale_each=False, pad_value=0, format=None, rotate=False):
    """Save a given Tensor into an image file.

    Args:
        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
            saves the tensor as a grid of images by calling ``make_grid``.
        fp (string or file object): A filename or a file object
        format(Optional):  If omitted, the format to use is determined from the filename extension.
            If a file object was used instead of a filename, this parameter should always be used.
        **kwargs: Other arguments are documented in ``make_grid``.
    """
    from PIL import Image
    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                     normalize=normalize, range=range, scale_each=scale_each, rotate=rotate)
    grid = cv2.cvtColor(grid, cv2.COLOR_RGB2BGR)
    im = Image.fromarray(grid).convert('RGB')
    logger.debug(f'save_image: {fp}')
    im.save(fp, format=format)
    # cv2.imwrite(fp, grid)

def lerp(start, end, weight):
    return start + weight * (end - start)

def translate_using_latent(nets, args, x_src, y_trg_list, z_trg_list, psi, filename):
    N, C, H, W = x_src.shape
    latent_dim = z_trg_list[0].shape[1]
    x_concat = [x_src]
    masks = nets.fan.get_heatmap(x_src) if args.w_hpf > 0 else None

    for i, y_trg in enumerate(y_trg_list):
        z_many = fluid.dygraph.to_variable(np.random.randn(10000, latent_dim)).astype(np.float32)
        # z_many = torch.randn(10000, latent_dim).to(x_src.device)
        y_many = np.random.randn(10000).astype(np.int32)
        y_many.fill(y_trg[0].numpy()[0])
        y_many = fluid.dygraph.to_variable(y_many).astype('int32')

        # y_many = torch.LongTensor(10000).to(x_src.device).fill_(y_trg[0])
        s_many = nets.mapping_network(z_many, y_many)
        s_avg = fluid.layers.reduce_mean(s_many, dim=0, keep_dim=True)
        # s_avg = torch.mean(s_many, dim=0, keepdim=True)
        s_avg = fluid.dygraph.to_variable(np.repeat(s_avg.numpy(), N, axis=0))
        # s_avg = s_avg.repeat(N, 1)

        for z_trg in z_trg_list:
            s_trg = nets.mapping_network(z_trg, y_trg)
            s_trg = lerp(s_avg, s_trg, psi)
            x_fake = nets.generator(x_src, s_trg, masks=masks)
            x_concat += [x_fake]

    x_concat = fluid.layers.concat(x_concat, axis=0)
    save_images(x_concat, N, filename)


def translate_and_reconstruct(nets, args, x_src, y_src, x_ref, y_ref, filename):
    N, C, H, W = x_src.shape
    s_ref = nets.style_encoder(x_ref, y_ref)
    masks = nets.fan.get_heatmap(x_src) if args.w_hpf > 0 else None
    x_fake = nets.generator(x_src, s_ref, masks=masks)
    s_src = nets.style_encoder(x_src, y_src)
    masks = nets.fan.get_heatmap(x_fake) if args.w_hpf > 0 else None
    x_rec = nets.generator(x_fake, s_src, masks=masks)
    x_concat = [x_src, x_ref, x_fake, x_rec]
    x_concat = fluid.layers.concat(x_concat, axis=0)
    save_images(x_concat, N, filename)
    del x_concat


def debug_image(nets, args, inputs, step):
    x_src, y_src = fluid.dygraph.to_variable(inputs.x_src).astype('float32'), fluid.dygraph.to_variable(inputs.y_src).astype('int32')
    x_ref, y_ref = fluid.dygraph.to_variable(inputs.x_ref).astype('float32'), fluid.dygraph.to_variable(inputs.y_ref).astype('int32')

    for name, model in nets.items():
        model.eval()
    N = inputs.x_src.shape[0]

    # translate and reconstruct (reference-guided)
    filename = ospj(args.sample_dir, '%06d_cycle_consistency.jpg' % (step))
    translate_and_reconstruct(nets, args, x_src, y_src, x_ref, y_ref, filename)

    # latent-guided image synthesis
    y_trg_list = [fluid.dygraph.to_variable(np.array(y).repeat(N))
                  for y in range(min(args.num_domains, 5))]
    # logger.debug(f'y_trg_list: {str(y_trg_list)}')
    z_trg_list = fluid.dygraph.to_variable(np.random.randn(args.num_outs_per_domain, 1, args.latent_dim).astype(np.float32).repeat(N, axis=1))
    # z_trg_list = torch.randn(args.num_outs_per_domain, 1, args.latent_dim).repeat(1, N, 1).to(device)
    for psi in [0.5, 0.7, 1.0]:
        filename = ospj(args.sample_dir, '%06d_latent_psi_%.1f.jpg' % (step, psi))
        translate_using_latent(nets, args, x_src, y_trg_list, z_trg_list, psi, filename)

    # reference-guided image synthesis
    filename = ospj(args.sample_dir, '%06d_reference.jpg' % (step))
    translate_using_reference(nets, args, x_src, x_ref, y_ref, filename)

    for name, model in nets.items():
        model.train()



